#include "prefetch/mithril.hpp"

/**
 * \file
 * \ingroup cachePrefetch
 * \brief cas::Mithril implementation.
 */

namespace cas
{
    SLOG_COMPONENT_DEFINE("Mithril");

    Mithril::Mithril()
    {
        SLOG_FUNCTION_NOARGS();
        recordTable.set_eviction(1);   // 1=evict_front
        prefetchTable.set_eviction(1); // 1=evict_front
    }

    Mithril::~Mithril()
    {
        SLOG_FUNCTION_NOARGS();
        recordTable.clear();
        miningTable.clear();
        prefetchTable.clear();
    }

    void Mithril::init()
    {
        SLOG_FUNCTION_NOARGS();

        // check requirement
        SASSERT_MSG(minSupport > 0, "Invalid mithril's minSupport");
        SASSERT_MSG(maxSupport > 0, "Invalid mithril's maxSupport");
        SASSERT_MSG(recordSize > 0, "Invalid mithril's recordSize");
        SASSERT_MSG(miningSize > 0, "Invalid mithril's miningSize");
        SASSERT_MSG(prefetchColSize > 0, "Invalid mithril's prefetchColSize");
        SASSERT_MSG(prefetchRowSize > 0, "Invalid mithril's prefetchRowSize");

        // allocate memory
        recordTable.clear();
        recordTable.resize(recordSize);
        miningTable.clear();
        prefetchTable.clear();
        prefetchTable.resize(prefetchRowSize);

        stat = MithrilStatistic{}; // clear stat
        traceChronoTime = 1;       // set timestamp = 1 to indicates started

        printf("[mithril] initializing with setting:\n");
        printf("[mithril] MIN_SUPPORT %u\tMAX_SUPPORT %u\tLOOKAHEAD_RANGE %u\n", minSupport, maxSupport, lookAheadRange);
        printf("[mithril] REC_SIZE %lu\tMINE_SIZE %lu\n", recordSize, miningSize);
        printf("[mithril] PREFETCH_ROW %lu\tPREFETCH_COL %lu\tPREFETCH_MEM %lu\n", prefetchRowSize, prefetchColSize,
               prefetchRowSize * (prefetchColSize + 1) * (1 << 9));
    }

    void Mithril::setMiningSize(const size_t &size)
    {
        if (traceChronoTime == 0)
        {
            SASSERT(size > 0);
            miningSize = size;
        }
        else
        {
            SFATAL_ERROR_CONT("Set miningSize before init");
        }
    }

    void Mithril::setLookAheadRange(const uint16_t &range)
    {
        SLOG_FUNCTION(range);
        lookAheadRange = range; // maybe dynamic
    }

    void Mithril::setParamSupport(const uint8_t &minSupport, const uint8_t &maxSupport)
    {
        // maybe dynamic
        if (minSupport)
            this->minSupport = minSupport;
        if (maxSupport)
            this->maxSupport = maxSupport;
    }

    void Mithril::setPrefetchSize(const uint16_t &col, const size_t &row)
    {
        if (traceChronoTime == 0)
        {
            SASSERT(col > 0 && row > 0);
            prefetchColSize = col;
            prefetchRowSize = row;
        }
        else
        {
            SFATAL_ERROR_CONT("Set prefetchSize before init");
        }
    }

    void Mithril::setRecordSize(const size_t &size)
    {
        if (traceChronoTime == 0)
        {
            SASSERT(size > 0);
            recordSize = size;
        }
        else
        {
            SFATAL_ERROR_CONT("Set recordSize before init");
        }
    }

    void Mithril::record(const CacheRequest &request)
    {
        chrono_clock_t stm = chrono_now();
        CachePage page = request.page;                    // copy of cache page
        const mithril_timestamp_t ts = traceChronoTime++; // timestamp start from 1, atomic
        stat.recordCount++;

        SASSERT(page.lba > 0);
        for (uint16_t i = 0; i < request.size; i++, page.lba++)
        {
            auto itm = miningTable.find(page);
            if (itm != miningTable.end())
            {
                // already in mining table
                if ((int32_t)itm->second.size() >= maxSupport)
                {
                    // evict for too frequent row
                    miningTable.erase(page);
                    stat.evictMiningTable++;
                }
                else
                {
                    // push mining table
                    itm->second.push_back(ts); // append timestamp
                    stat.recordMiningTable++;
                }
                continue;
            }

            if (minSupport != 1)
            {
                auto itr = recordTable.find(page);
                if (itr != recordTable.map_end())
                {
                    // already in record table, append to row
                    auto arr = recordTable.get(itr->second);
                    arr->push_back(ts); // append timestamp
                    stat.recordRecordTable++;

                    if ((int32_t)arr->size() == minSupport)
                    {
                        // reach minSupport, move to mining table
                        moveMiningTable(page, *arr, stm);
                        recordTable.erase(itr->second); // delete from record table
                    }
                }
                else
                {
                    // create new row, overwrite least recent if full
                    bool isEvict = recordTable.push_back(page, std::vector<mithril_timestamp_t>{ts});
                    stat.evictRecordTable += isEvict;
                    stat.recordRecordTable++;
                    stat.lineRecordTable++;
                }
            }
            else
            {
                // minSupport == 1, direct insert to mining table
                moveMiningTable(page, std::vector<mithril_timestamp_t>{ts}, stm);
            }
        }
        stat.timeTotalRecord += chrono_now() - stm; // timing
    }

    void Mithril::mining()
    {
        SLOG_FUNCTION_NOARGS();
        const chrono_clock_t stm = chrono_now();

        // transform and sort
        std::vector<MiningTableRow> tab(miningTable.begin(), miningTable.end());
        std::sort(tab.begin(), tab.end(), [](const MiningTableRow &_a, const MiningTableRow &_b)
                  { return _a.second[0] < _b.second[0]; });

        const size_t tlen = tab.size(); // tab length
        for (size_t i = 0; i < tlen; i++)
        {
            const mithril_timestamp_t bts = tab[i].second[0];
            const size_t blen = tab[i].second.size();
            for (size_t j = i + 1; j < tlen; j++)
            {
                const mithril_timestamp_t ots = tab[j].second[0];
                SASSERT(ots >= bts);
                if (ots - bts > lookAheadRange)
                {
                    break; // beyond look-ahead range
                }
                if (tab[j].second.size() == blen)
                {
                    // check association
                    uint8_t assoc = checkAssociation(tab[i].second, tab[j].second);
                    if (assoc > 0)
                    {
                        pushAssociation(tab[i].first, tab[j].first); // add new association
                    }
                }
            }
        }

        // clear up
        miningTable.clear();
        tab.clear(); // for safety

        stat.timeTotalMining += chrono_now() - stm; // timing
    }

    uint8_t Mithril::checkAssociation(const std::vector<mithril_timestamp_t> &a, const std::vector<mithril_timestamp_t> &b)
    {
        /* check distance between each trace pair, for all 1 = strong assoc. */
        bool isStrong = true;
        for (auto ia = a.begin(), ib = b.begin(); ia != a.end(); ia++, ib++)
        {
            const auto dist = (*ia > *ib) ? (*ia - *ib) : (*ib - *ia);
            if (dist > lookAheadRange)
            {
                return 0; // beyond look-ahead range
            }
            isStrong |= dist <= 1;
        }
        return isStrong ? 2 : 1;
    }

    void Mithril::pushAssociation(const CachePage &a, const CachePage &b)
    {
        auto ita = prefetchTable.find(a);
        if (ita != prefetchTable.map_end())
        {
            // located page `a`
            auto arr = prefetchTable.get(ita->second);
            bool isEvict = arr->size() >= prefetchColSize; // reach max length
            if (isEvict)
            {
                arr->erase(arr->begin());
            }
            arr->push_back(b);
            prefetchTable.move_back(ita->second); // update as most recent

            stat.evictPrefetchList += isEvict;
            stat.recordPrefetchTable++;
        }
        else
        {
            // page `a` not found, add a new
            if (prefetchTable.capacity() == prefetchTable.size())
            {
                // full, evict least recent
                auto arr = prefetchTable.pop_front().second;
                stat.evictPrefetchList += arr.size();
                stat.evictPrefetchTable++;
                arr.clear(); // free memory
            }
            SASSERT(false == prefetchTable.push_back(a, std::vector<CachePage>{b})); // should not evict
            stat.recordPrefetchTable++;
            stat.linePrefetchTable++;
        }
    }

    void Mithril::moveMiningTable(const CachePage &page, const std::vector<mithril_timestamp_t> &arr, chrono_clock_t &stm)
    {
        SASSERT(miningTable.size() < miningSize);
        miningTable.emplace(page, arr);
        stat.recordMiningTable += minSupport;
        stat.lineMiningTable++;

        // if mining table full, start mining
        if (miningTable.size() == miningSize)
        {
            stat.timeTotalRecord += chrono_now() - stm; // timing
            mining();
            // \note To async job, use std::async
            // std::async(std::launch::async, []{ mining(); });
            stat.miningCount++;
            stm = chrono_now(); // reset timer
        }
    }

    void Mithril::updateAssociation(const CachePage &a, const CachePage &b, int32_t scoreOffset)
    {
        auto it = prefetchTable.find(a);
        if (it == prefetchTable.map_end())
        {
            return;
        }
        // append items
        auto arr = prefetchTable.get(it->second);
        auto itb = std::find(arr->begin(), arr->end(), b);
        if (scoreOffset > 0)
        {
            // stronger association
            if (itb == arr->end())
            {
                bool isEvict = arr->size() >= prefetchColSize;
                if (isEvict)
                {
                    arr->erase(arr->begin()); // reach max length
                }
                arr->push_back(b);
            }
            else
            {
                arr->erase(itb);
                arr->push_back(b);
            }
            prefetchTable.move_back(it->second); // update as most recent
        }
        else if (scoreOffset < 0)
        {
            // weaker association
            if (itb != arr->end())
            {
                arr->erase(itb);
                arr->insert(arr->begin(), b);
            }
        }
    }

    void Mithril::print() const
    {
        const auto &s = stat;
        printf("[mithril]          exec     size     records  lines    evict\n");
        printf("[mithril] record   %8lu %8lu %8lu %8lu %8lu\n",
               s.recordCount, recordTable.size(), s.recordRecordTable, s.lineRecordTable, s.evictRecordTable);
        printf("[mithril] mining   %8lu %8lu %8lu %8lu %8lu\n",
               s.miningCount, miningTable.size(), s.recordMiningTable, s.lineMiningTable, s.evictMiningTable);
        printf("[mithril] prefetch %8lu %8lu %8lu %8lu %8lu\n",
               s.prefetchCount, prefetchTable.size(), s.recordPrefetchTable, s.linePrefetchTable, s.evictPrefetchTable);
        printf("[mithril] time record %10lf mining %10lf prefetch %10lf\n",
               s.timeTotalRecord.count(), s.timeTotalMining.count(), s.timeTotalPrefetch.count());
        printf("[mithril] assoc unique %8lu query %8lu evict %8lu inplace %5lu\n",
               s.uniquePrefetchAssoc, s.queryPrefetchAssoc, s.evictPrefetchList, s.inplacePrefetchAssoc);
    }

    std::vector<CachePage> Mithril::prefetch(CacheRequest request)
    {
        const chrono_clock_t stm = chrono_now();
        CachePage page = request.page; // copy of cache page
        std::set<CachePage> ret;       // must be unique

        stat.prefetchCount++;

        for (uint16_t i = 0; i < request.size; i++, page.lba++)
        {
            auto it = prefetchTable.find(page);
            if (it != prefetchTable.map_end())
            {
                // append items
                auto arr = prefetchTable.get(it->second);
                stat.queryPrefetchAssoc += arr->size();
                ret.insert(arr->begin(), arr->end());
            }
        }

        page.lba -= request.size;
        for (uint16_t i = 0; i < request.size; i++, page.lba++)
        {
            // erase page already in request
            auto cnt = ret.erase(page);
            stat.inplacePrefetchAssoc += cnt;
        }

        stat.uniquePrefetchAssoc += ret.size();
        stat.timeTotalPrefetch += chrono_now() - stm; // timing
        return std::vector<CachePage>(ret.begin(), ret.end());
    }

} // namespace cas
